#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, find its orientation and
## extract the rectified.
##
## Changes that must be performed as soon as possible: read the
## intrinsic parameters from somewhere, and make it easy to switch to
## the equirectangular model.

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0] == 'extract_edgels-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io

from numpy import dot

# from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors
from camori_harris import PictureHarris

from quaternion import Quat, random_quaternion

import simplejson

import scipy.optimize
#from fmin_sa import Sampling
from simplex_sa import SimplexSO3

import Image

from plot_aux import *

set_printoptions(precision=3)

if __name__ == '__main__':

    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0] == 'extract_edgels-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    #################################################################
    ## Load image and initialize pic object

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print sys.argv[0], '<job_file.json> <frame_number>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    framenum = int(sys.argv[2])
    filename = fileroot+'/frames/'+job_params['filename_format']%framenum

    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel

    # Smooth out
    if ("gaussian_smoothing_factor" in job_params.keys() and 
        job_params["gaussian_smoothing_factor"] > 0):
        for c in range(3):
            imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], \
                          double(job_params["gaussian_smoothing_factor"]))

    ## Creates picture object
    # if not 'projection_model' in job_params:
    #     model = 'perspective'
    # else:
    #     model = job_params['projection_model']

    # if model != 'harris':
    #     raise NotImplementedError
    # elif model == 'harris':
    if True:
        ## Intrinsic parameters
        p_point = array(job_params['principal_point_guess'])
        focal_distance = job_params['focal_distance_guess']
        distortion = array(job_params['distortion_coefficient_guess'])
        distortion = -300
        pic = PictureHarris(imr, p_point[0], p_point[1], focal_distance, distortion)


    ## Run the grid sweep, and generate the edgels list.
    ## Edgel extractor parameters
    gspc = job_params['grid_spacing']
    glim = job_params['gradient_threshold'] # ** -1
    pic.extract_edgels(gspc, glim)
    ##
    #################################################################



    def loss_ori(qq, *params):
        loss_par = (1/.15,3,1.0)
        q = qq
        kappa = params[3]
        return pic.fit_error(params[0], params[1], params[2], kappa,
                      q, *loss_par)

    xini = array([p_point[0], p_point[1], focal_distance, distortion])
    ss = SimplexSO3(loss_ori, Quat(0,.36,0), xini)

    #################################################################
    ## Perform global optimization procedure over all the intrinsic
    ## parameters and orientation constrained to the 3-sphere.
    def loss_cal(pp, ss, ori_tol):
        print 70*'-'
        print pp
        ss.params = pp
        #ss.reset_state(ss.best)
        ss.best_reset()
        ss.opt_loop(10000, xtol=ori_tol, ftol=ori_tol)
        print ss.best
        return ss.fbest



    xopt = scipy.optimize.fmin(loss_cal, xini, args=(ss,1e-5), xtol=1e-5, ftol=1e-5, disp=0)
    print '--', xopt

    xopt = scipy.optimize.fmin(loss_cal, xopt, args=(ss,1e-8), xtol=1e-10, ftol=1e-10, disp=0)
    print '--', xopt

    print 'loss', loss_ori(ss.best, *xopt)

    #print xopt

    pic.middlex = xopt[0]
    pic.middley = xopt[1]
    pic.fd = xopt[2]
    pic.distortion = xopt[3]

    qopt = Quat(ss.solution()).canonical()
    print 'pp = ', (pic.middlex, pic.middley)
    print 'fd = ', pic.fd
    print 'kappa = ', pic.distortion
    print 'qopt',qopt
    
    ##
    #################################################################


    plot_problem(pic, 0)

    figure(11)
    imshow(pic.frame*.9/255.)
    pic.plot_vdirs(gca(), pic.Iwidth/40.1, qopt)
    axis('equal')
    grid()

        # figure(10)
        # q = Quat(ss.samples[argmin(ss.mlogL)])
        # imshow(.1+pic.frame*.9/255.)
        # pic.plot_vdirs(gca(), 30, q)
        # axis('equal')
        # grid()
